import qualified Data.Set as Set
import Data.List (findIndices)

type Operand = Int
data Operation = Nop | Acc | Jmp deriving (Show, Eq)
data Instruction = Instruction Operation Operand deriving (Show)

type PC = Int
type RegisterValue = Int
type RunRecord = Set.Set PC
type State = (PC, RegisterValue, RunRecord)

parseOperation :: String -> Operation
parseOperation string = case string of
    "nop" -> Nop
    "acc" -> Acc
    "jmp" -> Jmp
    _ -> error "Unkown operation" 

parseInstruction :: String -> Instruction
parseInstruction string = 
    Instruction (parseOperation operationString) (sign * read operandString)
    where
        operationString = take 3 string
        sign = if (string !! 4) == '+' then 1 else -1
        operandString = drop 5 string

parseInput :: String -> [Instruction]
parseInput = map parseInstruction . lines

runInstruction :: Instruction -> State -> State
runInstruction (Instruction operation operand) (pc,regValue,record) = 
    case operation of 
        Nop -> (pc+1, regValue, newRecord)
        Acc -> (pc+1, regValue+operand, newRecord)
        Jmp -> (pc+operand, regValue, newRecord) 
    where
        newRecord = Set.insert pc record

run :: [Instruction] -> Maybe RegisterValue
run insts = doRun insts (0,0,Set.empty)
    where
        doRun :: [Instruction] -> State -> Maybe RegisterValue
        doRun insts state@(pc,regValue,record)
            | Set.member pc record = Nothing
            | pc >= (length insts) = Just regValue
            | otherwise = 
                doRun insts newState
                where
                    inst = insts !! pc
                    newState = runInstruction inst state

changeOperation :: Operation -> [Instruction] -> PC -> [Instruction]
changeOperation operation ((Instruction _ operand):insts) 0 = ((Instruction operation operand):insts)
changeOperation operation (inst:insts) pc =  
    if pc > 0 
        then inst : (changeOperation operation insts (pc-1))
        else error "PC too bigg!"

getOperation :: Instruction -> Operation
getOperation (Instruction operation _) = operation

otherOperation :: Operation -> Operation
otherOperation oper = case oper of 
    Jmp -> Nop
    Nop -> Jmp

try :: [Instruction] -> RegisterValue 
try insts =  unwrap 
           . head 
           . filter (/=Nothing)
           . map run $ tryInstsLst
    where
        getNewInstsLst :: Operation -> [[Instruction]]
        getNewInstsLst operation = map (changeOperation (otherOperation operation) insts) 
                                 . findIndices (\inst -> (getOperation inst) == operation) $ insts
        tryInstsLst = (getNewInstsLst Jmp) ++ (getNewInstsLst Nop)

isNothig :: Maybe a -> Bool
isNothig Nothing = True
isNothig _ = False

unwrap :: Maybe a -> a
unwrap (Just v) = v

main :: IO ()
main = do
    input <- readFile "./input.txt"     
    print . try . parseInput $ input
